pub const MILLAR_RABIN_WITNESSES: &[u128] = &[2, 325, 9375, 28178, 450775, 9780504, 1795265022];

pub fn millar_rabin(n: u128) -> bool {
    match n {
        0 | 1 => return false,
        2 | 3 | 5 => return true,
        _ if n % 2 == 0 || n % 3 == 0 || n % 5 == 0 => return false,
        _ => {}
    }
    // n - 1 = 2^t * u
    let t = (n - 1).trailing_zeros();
    let u = (n - 1) >> t;
    'witness: for &witness in MILLAR_RABIN_WITNESSES {
        let a = witness % n;
        if a == 0 || a == n - 1 {
            continue;
        }
        let mut x = pow(a, u, n);
        if x == 1 {
            continue;
        }
        for _ in 0..t {
            if x == n - 1 {
                continue 'witness;
            }
            x = mul(x, x, n);
        }
        return false;
    }
    true
}

fn pow(a: u128, mut b: u128, n: u128) -> u128 {
    let mut ans = 1_u128;
    let mut a = a % n;
    while b > 0 {
        if b % 2 == 1 {
            ans = mul(ans, a, n);
        }
        a = mul(a, a, n);
        b /= 2;
    }
    ans
}

fn mul(a: u128, b: u128, n: u128) -> u128 {
    if n < 1 << 64 {
        return (a % n) * (b % n) % n;
    }
    let mut ans = 0_u128;
    let mut a = a % n;
    let mut b = b % n;
    while b > 0 {
        if b % 2 == 1 {
            ans = (ans + a) % n;
        }
        a = (a + a) % n;
        b /= 2;
    }
    ans
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_millar_rabin() {
        for prime in [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31] {
            assert!(millar_rabin(dbg!(prime)));
        }
        for composite in [4, 6, 8, 9, 10, 12, 14, 15, 16, 18, 20] {
            assert!(!millar_rabin(composite));
        }
    }
}
